#include "socks5.h"
#include "LockAndUnlock.h"

void CSocks5Server::_DoConnectEventHandle(int connectfd)
{
	Connect* connect = new Connect;
	connect->_clientChannel._fd = connectfd;
	connect->_clientChannel._event = EPOLLIN;
	connect->_status = AUTH;
	connect->_counter++;
	_connectMap[connectfd] = connect;

	SetNonblocking(connectfd);
	OperateEvent(connectfd, EPOLLIN, EPOLL_CTL_ADD);
}

// 读取客户端发来的协商版本和认证方法的请求
int CSocks5Server::Auth(int connectfd)
{
	TraceLog("Auth : %d", connectfd);
	char buf[258];
	int recv_length = recv(connectfd, buf, 258, MSG_PEEK);
	if(recv_length < 0){
		ErrorLog("read socks5 auth head");
		return -1;
	}

	if(recv_length < 3){
		return 0;
	}
  // 读取客户端发来的请求, 协商版本和认证方法
  /*
   * + ---- + -------- + --------- +
   * | VER  | NMETHODS |  METHODS  |
   * + ---- + -------- + --------- +
   * |  1   |    1     |  1 to 255 |
   * + ---- + -------- + --------- +
   */
	recv(connectfd, buf, recv_length, 0);
	Decrypt(buf, recv_length);
	TraceLog("recv_length:%d", recv_length);

  // VER字段代表 Socks 的版本, 此处应为被设置成 0x05
	if (buf[0] != 0x05){
		ErrorLog("The VeR not socks5");
		return -1;
	}
	return 1;
}

// 上一步的客户端协商认证环节结束后, 客户端就发送详细的请求信息
int CSocks5Server::EstablishConnection(int connectfd)
{
	TraceLog("%d", connectfd);

	const size_t len = 256;
	char buf[len];

	// 1. 先进行窥探,查看缓冲区是否有足够数据
	int recv_length = recv(connectfd, buf, len, MSG_PEEK);
	// 1.1 小于 0, 出错
	if (recv_length < 0){
		ErrorLog("read request");
		return -1;
	}
	// 1.2 等于 0, 发送方已断开连接
	else if (0 == recv_length){
		TraceLog("recv 0");
		RemoveConnect(connectfd);
		return -1;
	}

	// 1.3 收到数据数量小于 10
	if (recv_length < 10){
		TraceLog("recv len: %d", recv_length);
		return 0;
	}

  // 2. 先接受请求的前 4 个字节, 并解密
	recv(connectfd, buf, 4, 0);
	Decrypt(buf, 4);
  /*   版本          保留 后面地址类型 目的地址  以网络字节顺序出现的端口号
   * + ---- + ---- + ---- + ------ + -------- + -------- +
   * | VER  | CMD  | RSV  |  ATYP  | DST.ADDR | DST.port |
   * + ---- + ---- + ---- + ------ + -------- + -------- +
   * |  1   |  1   | 0x00 |    1   | Variable |    2     |
   * + ---- + ----- + --- + ------ + -------- + -------- +
   */
	char ip[16];
	char port[2];
  // 3. 根据 ATYP 字段判断 DST.ADDR 字段的具体内容
  // 3.1 0x01 : IPv4
	if (buf[3] == 0x01){
		if(4 != recv(connectfd, ip, 4, 0)){
			ErrorLog("recv IPv4 error!!!\n");
			return -1;
		}
		ip[4] = '\0';
		Decrypt(buf, 4);
		TraceLog("recv ipv4 : %s", ip);
	}
  // 3.2  0x03 : 域名, 如果是域名情况, 它的 DST.ADDR 部分内容如下:
  //      第一个字节为域名长度, 剩余的内容为域名，没有'\0'结尾
	else if (buf[3] == 0x03){
		if(1 != recv(connectfd, buf, 1, 0)){
			ErrorLog("recv domain name len!!!");
			return -1;
		}

		Decrypt(buf, 1);
		int len = buf[0]; 
		if(len != recv(connectfd, buf, len, 0)){
			ErrorLog("get domain name fail!!!");
			return -1;
		}

		buf[len] = '\0';
		//TraceLog("encty domain:%s", buf);

		Decrypt(buf, len);
		TraceLog("decrypt domain : %s", buf);

		// 通过域名取ip
    // gethostbyname() : 用域名或者主机名获取地址, 操作系统提供的库函数
		struct hostent* hptr = gethostbyname(buf);

    // 失败返回 NULL 
		if(NULL == hptr){
			ErrorLog("gethostbyname() to %s fail !!!",buf);
			return -1;
		}

    // 成功返回指向 hostent 结构体的指针
    // struct hostent {
    //   char  *h_name;            /* official name of host */
    //   char **h_aliases;         /* alias list */
    //   int    h_addrtype;        /* host address type */
    //   int    h_length;          /* length of address */
    //   char **h_addr_list;       /* list of addresses */
    // }
		struct in_addr addr;
    // # define    h_addr  h_addr_list[0]
		memcpy(ip, hptr->h_addr, hptr->h_length);
	}
  // 3.3 0x04 : IPv6 此处不处理,直接报错
	else if (buf[3] == 0x04){
		TraceLog("Do't support ipv6");
		return -1;
	}
	else{
		ErrorLog("invalid Address Type");
	}

	// 4. 获取端口号 
	if(2 != recv(connectfd, port, 2, 0)){
		ErrorLog("read port");
		return -1;
	}
	Decrypt(port, 2);

	struct sockaddr_in addr;
	addr.sin_family = AF_INET;
	memcpy(&addr.sin_addr.s_addr, ip, 4);
	addr.sin_port = *((uint16_t*)port);

	int serverfd = socket(AF_INET, SOCK_STREAM, 0);
	if(serverfd < 0){
		ErrorLog("socket");
		return -1;
	}

	if (connect(serverfd, (struct sockaddr*)&addr, sizeof(addr)) < 0){
		ErrorLog("connect error");
		close(serverfd);
		return -1;
	}

	SetNonblocking(serverfd);
	return serverfd;
}

// 处理客户端的请求(针对不同状态给予不同回复), 如果是转发状态则直接建立连接转发数据 
void CSocks5Server::_DoReadEventHandle(int connectfd)
{
  std::map<int, Connect*>::iterator it = _connectMap.find(connectfd);
	if(it != _connectMap.end()){
		Connect* connect = it->second;
    // 1. 服务器收到客户端的协商验证请求后, 从客户端给定的方法中选择一个
    //    并发送一个方法选中的消息回客户端
		if(AUTH == connect->_status){
			char response[2];
			response[0]= 0x05;
			int ret = Auth(connectfd);
      // 1.1 回应0x00 代表不需要验证
			if(ret == 1){
				response[1] = 0x00;
				connect->_status = ESTABLISHED;
			}
      // 1.2 回应 0xFF 代表没有可以接受的方法
      //     发送 0xFF 让客户端关闭连接
			else if(-1 == ret){
				response[1] = 0xFF;
				RemoveConnect(connectfd);
			}
			else{
				return;
			}
			// 1.3 回复客户端,先对数据进行加密
			Encrypt(response, 2);
			if(2 != send(connectfd, response, 2, MSG_DONTWAIT)){
				ErrorLog("auth response");
			}
		} // end if(AUTH == connect->_status)

    // 2. 如果已经完成了建立了一个到SOCKS服务器的连接, 并且完成了认证方式的协商过程, 
    //    客户机将会发送一个SOCKS请求信息给服务器. 服务器将会根据请求返回如下格式响应.
		else if(connect->_status == ESTABLISHED){
			char response[10] = {0};
			response[0] = 0x05; // 协议版本 socks5
			response[3] = 0x01; // 地址类型 0x01 代表 IPv4

			int serverfd = EstablishConnection(connectfd);
			if (0 == serverfd){
				return;
			}
      // 2.1 0x01 代表普通的SOCKS服务器请求失败
			else if (serverfd == -1){
				response[2] = 0x01;
			}
      // 2.2 0x00 代表成功, 准备进行数据转发
			else{
				response[1] = 0x00;
				connect->_status = FORWARD;
			}
      // 建立连接
			connect->_serverChannel._fd = serverfd;
			connect->_counter++;
			OperateEvent(serverfd, EPOLLIN, EPOLL_CTL_ADD);
			_connectMap[serverfd] = connect;

      // 2.3 发送回应给客户端
			Encrypt(response, 10);
			if (10 != send(connectfd, response, 10, MSG_DONTWAIT)){
				ErrorLog("server send response to client");
			}
		} // end if(connect->_status == ESTABLISHED)

    // 3. 服务器此时已经进入转发状态
		else if (connect->_status == FORWARD){
			Channel* clientChannel = &connect->_clientChannel;
			Channel* serverChannel = &connect->_serverChannel;
			bool recvDecrypt = true;
			bool sendEncry = false;

			if (connectfd == serverChannel->_fd){
				std::swap(clientChannel, serverChannel);
				std::swap(recvDecrypt, sendEncry);
			}
			Forward(clientChannel, serverChannel, recvDecrypt, sendEncry);
		}
	} // end if(it != _connectMap.end()){

	else{
		assert(false);
	}
}

int main()
{
	CSocks5Server server(8800);
	server.StartUp();
}
